package org.wang.generalcode.modules.knowledge.vector.service;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.*;
import io.milvus.v2.service.index.request.CreateIndexReq;
import io.milvus.v2.service.vector.request.DeleteReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.QueryReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentMetadata;
import org.springframework.ai.ollama.OllamaEmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.stereotype.Service;
import org.wang.generalcode.modules.knowledge.vector.config.MilvusArchive;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author xtwang
 * @des MilvusServiceImpl
 * @date 2025/2/25 下午3:10
 */
@Service
@RequiredArgsConstructor
public class MilvusServiceImpl implements MilvusService {

	private final MilvusClientV2 milvusClientV2;
	private final OllamaEmbeddingModel ollamaEmbeddingModel;

	/**
	 * 检查集合是否存在
	 */
	@Override
	public void hasCollection() {
		// milvusClientV2.dropCollection(DropCollectionReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build());
		Boolean b = milvusClientV2
				.hasCollection(HasCollectionReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build());
		if (!b) {
			this.createCollection();
		}
	}

	/**
	 * 创建集合
	 */
	public void createCollection() {
		// 创建字段
		CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema()
				// 创建主键字段
				.addField(AddFieldReq.builder()
						// 字段名
						.fieldName(MilvusArchive.Field.ID)
						// 字段描述
						.description("主键ID")
						// 字段类型
						.dataType(DataType.Int64)
						// 是否为主键
						.isPrimaryKey(true)
						// 设置主键自增
						.autoID(true)
						.build())
				.addField(AddFieldReq.builder()
						// 字段名
						.fieldName(MilvusArchive.Field.FILE_NAME)
						// 字段描述
						.description("文件名")
						// 字段类型
						.dataType(DataType.VarChar)
						// 设置字段为可空
						.isNullable(true)
						.build())
				// 创建特征向量字段
				.addField(AddFieldReq.builder()
						// 字段名
						.fieldName(MilvusArchive.Field.FEATURE)
						// 字段描述
						.description("特征向量")
						// 字段类型
						.dataType(DataType.FloatVector)
						// 设置向量维度
						.dimension(MilvusArchive.FEATURE_DIM)
						.build())
				.addField(AddFieldReq.builder()
						// 字段名
						.fieldName(MilvusArchive.Field.TEXT)
						// 字段描述
						.description("文本")
						// 字段类型
						.dataType(DataType.VarChar)
						// 设置字段为可空
						.isNullable(true)
						.build())
				.addField(AddFieldReq.builder()
						// 字段名
						.fieldName(MilvusArchive.Field.METADATA)
						// 字段描述
						.description("元数据")
						// 字段类型
						.dataType(DataType.VarChar)
						// 设置字段为可空
						.isNullable(true)
						.build());
		// 创建集合
		CreateCollectionReq collectionReq = CreateCollectionReq.builder()
				// 集合名称
				.collectionName(MilvusArchive.COLLECTION_NAME)
				// 集合描述
				.description("自定义知识库")
				// 集合字段
				.collectionSchema(schema)
				// 分片数量
				.numShards(MilvusArchive.SHARDS_NUM)
				.build();
		milvusClientV2.createCollection(collectionReq);

		// 创建索引
		IndexParam indexParam = IndexParam.builder()
				// 索引字段名
				.fieldName(MilvusArchive.Field.FEATURE)
				// 索引类型
				.indexType(IndexParam.IndexType.IVF_FLAT)
				// 索引距离度量
				.metricType(IndexParam.MetricType.COSINE)
				.build();
		CreateIndexReq createIndexReq = CreateIndexReq.builder()
				.collectionName(MilvusArchive.COLLECTION_NAME)
				.indexParams(Collections.singletonList(indexParam))
				.build();

		milvusClientV2.createIndex(createIndexReq);
	}

	/**
	 * 插入数据
	 */
	@Override
	public InsertResp insert(float[] vectorParam, String text, String metadata, String fileName) {
		// 校验集合是否存在
		this.hasCollection();
		JsonObject jsonObject = new JsonObject();
		// 数组转换成JsonElement
		jsonObject.add(MilvusArchive.Field.FEATURE, new Gson().toJsonTree(vectorParam));
		jsonObject.add(MilvusArchive.Field.TEXT, new Gson().toJsonTree(text));
		jsonObject.add(MilvusArchive.Field.METADATA, new Gson().toJsonTree(metadata));
		jsonObject.add(MilvusArchive.Field.FILE_NAME, new Gson().toJsonTree(fileName));
		InsertReq insertReq = InsertReq.builder()
				// 集合名称
				.collectionName(MilvusArchive.COLLECTION_NAME)
				.data(Collections.singletonList(jsonObject))
				.build();

		return milvusClientV2.insert(insertReq);
	}

	/**
	 * 批量插入数据
	 *
	 * @param vectorParam 向量参数
	 * @param text        文本
	 * @param metadata    元数据
	 * @param fileName    文件名
	 */
	@Override
	public InsertResp batchInsert(List<float[]> vectorParam, List<String> text, List<String> metadata,
			List<String> fileName) {
		if (vectorParam.size() == text.size() && vectorParam.size() == metadata.size()
				&& vectorParam.size() == fileName.size()) {
			List<JsonObject> jsonObjects = new ArrayList<>();
			for (int i = 0; i < vectorParam.size(); i++) {
				JsonObject jsonObject = new JsonObject();
				// 数组转换成JsonElement
				jsonObject.add(MilvusArchive.Field.FEATURE, new Gson().toJsonTree(vectorParam.get(i)));
				jsonObject.add(MilvusArchive.Field.TEXT, new Gson().toJsonTree(text.get(i)));
				jsonObject.add(MilvusArchive.Field.METADATA, new Gson().toJsonTree(metadata.get(i)));
				jsonObject.add(MilvusArchive.Field.FILE_NAME, new Gson().toJsonTree(fileName.get(i)));
				jsonObjects.add(jsonObject);
			}
			InsertReq insertReq = InsertReq.builder()
					// 集合名称
					.collectionName(MilvusArchive.COLLECTION_NAME)
					.data(jsonObjects)
					.build();
			return milvusClientV2.insert(insertReq);
		}
		return null;
	}

	/**
	 * 搜索数据
	 *
	 * @param vectorParam 向量参数
	 */
	@Override
	public SearchResp search(float[] vectorParam) {
		this.loadCollection();
		FloatVec floatVec = new FloatVec(vectorParam);
		SearchReq searchReq = SearchReq.builder()
				// 集合名称
				.collectionName(MilvusArchive.COLLECTION_NAME)
				// 搜索距离度量
				.metricType(IndexParam.MetricType.COSINE)
				// 搜索向量
				.data(Collections.singletonList(floatVec))
				// 搜索字段
				.annsField(MilvusArchive.Field.FEATURE)
				// 返回字段
				.outputFields(Arrays.asList(MilvusArchive.Field.ID, MilvusArchive.Field.TEXT,
						MilvusArchive.Field.METADATA, MilvusArchive.Field.FILE_NAME))
				// 搜索数量
				.topK(5)
				.build();
		return milvusClientV2.search(searchReq);
	}

	@PostConstruct
	public void loadCollection() {
		// 查询集合状态
		if (!milvusClientV2
				.getLoadState(GetLoadStateReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build())) {
			milvusClientV2
					.loadCollection(LoadCollectionReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build());
		}
	}

	public void releaseCollection() {
		milvusClientV2.releaseCollection(
				ReleaseCollectionReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).build());
	}

	/**
	 * 删除数据
	 *
	 * @param ids id
	 */
	@Override
	public void delete(String... ids) {
		milvusClientV2.delete(
				DeleteReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).ids(Arrays.asList(ids)).build());
	}

	@Override
	public void delete(String filter) {
		milvusClientV2.delete(DeleteReq.builder().collectionName(MilvusArchive.COLLECTION_NAME).filter(filter).build());
	}

	/**
	 * 搜索数据
	 *
	 * @param request 搜索请求
	 */
	@Override
	public List<Document> search(SearchRequest request) {
		String query = request.getQuery();
		float[] embed = ollamaEmbeddingModel.embed(query);
		List<List<SearchResp.SearchResult>> searchResults = this.search(embed).getSearchResults();
		if (!searchResults.isEmpty()) {
			return searchResults.get(0).stream()
					.filter(searchResult -> searchResult.getScore() >= request.getSimilarityThreshold())
					.map(searchResult -> {
						String docId = searchResult.getId().toString();
						String content = searchResult.getEntity().get(MilvusArchive.Field.TEXT).toString();
						JSONObject metadata = JSON
								.parseObject(searchResult.getEntity().get(MilvusArchive.Field.METADATA).toString());
						metadata.put(DocumentMetadata.DISTANCE.value(), 1.0F - searchResult.getScore());
						return Document.builder().id(docId).text(content).metadata(metadata.getInnerMap())
								.score((double) searchResult.getScore()).build();
					}).toList();
		}
		return List.of();
	}

	/**
	 * 查询数据
	 *
	 * @param fileName 文件名模糊搜索
	 */
	@Override
	public IPage<List<QueryResp.QueryResult>> query(IPage iPage, String fileName) {
		// 查询集合状态并加载到内存
		this.loadCollection();
		// 先查询总数
		QueryReq countReq = QueryReq.builder()
				.collectionName(MilvusArchive.COLLECTION_NAME)
				.outputFields(List.of(MilvusArchive.Field.ID))
				.filter(StringUtils.isNotBlank(fileName)
						? String.format("%s LIKE \"%%%s%%\"", MilvusArchive.Field.FILE_NAME, fileName)
						: String.format("%s > 0", MilvusArchive.Field.ID))
				.build();
		long total = milvusClientV2.query(countReq).getQueryResults().size();

		// 再查询分页数据
		QueryReq queryReq = QueryReq.builder()
				.collectionName(MilvusArchive.COLLECTION_NAME)
				.offset((long) (iPage.getCurrent() - 1) * iPage.getSize())
				.limit(iPage.getSize())
				.outputFields(List.of(MilvusArchive.Field.FILE_NAME, MilvusArchive.Field.TEXT))
				.filter(StringUtils.isNotBlank(fileName)
						? String.format("%s LIKE \"%%%s%%\"", MilvusArchive.Field.FILE_NAME, fileName)
						: String.format("%s > 0", MilvusArchive.Field.ID))
				.build();
		QueryResp query = milvusClientV2.query(queryReq);
		iPage.setRecords(query.getQueryResults());
		iPage.setTotal(total);
		return iPage;
	}
}
